import torch
from torchvision import datasets, transforms
import matplotlib.pyplot as plt
import numpy as np
import os

def load_mnist_dataset():
    """加载MNIST数据集"""
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])
    
    # 加载训练集
    try:
        train_dataset = datasets.MNIST('./data', train=True, download=True, transform=transform)
        test_dataset = datasets.MNIST('./data', train=False, download=True, transform=transform)
        return train_dataset, test_dataset
    except Exception as e:
        print(f"加载数据集时出错: {e}")
        return None, None

def visualize_mnist_samples(dataset, num_samples=25, rows=5, cols=5, save_path='mnist_samples.png'):
    """可视化MNIST数据集的样本"""
    if dataset is None:
        print("数据集未加载，无法可视化")
        return
    
    # 创建图形
    plt.figure(figsize=(12, 12))
    
    # 随机选择样本
    indices = np.random.choice(len(dataset), num_samples, replace=False)
    
    # 绘制样本
    for i, idx in enumerate(indices):
        if i >= rows * cols:
            break
            
        # 获取图像和标签
        image, label = dataset[idx]
        
        # 将图像转换为numpy数组并反归一化
        image = image.numpy()[0]  # 取第一个通道
        
        # 绘制子图
        plt.subplot(rows, cols, i + 1)
        plt.imshow(image, cmap='gray')
        plt.title(f"标签: {label}")
        plt.axis('off')
    
    plt.tight_layout()
    
    # 保存图像
    plt.savefig(save_path)
    print(f"样本可视化已保存为 '{save_path}'")
    
    # 显示图像
    plt.show()

def visualize_digit_samples(dataset, digit, num_samples=25, rows=5, cols=5, save_path=None):
    """可视化特定数字的样本"""
    if dataset is None:
        print("数据集未加载，无法可视化")
        return
    
    # 找到所有指定数字的索引
    digit_indices = []
    for i in range(len(dataset)):
        _, label = dataset[i]
        if label == digit:
            digit_indices.append(i)
            if len(digit_indices) >= num_samples:
                break
    
    # 如果找不到足够的样本
    if len(digit_indices) < num_samples:
        num_samples = len(digit_indices)
        rows = int(np.ceil(np.sqrt(num_samples)))
        cols = rows
    
    # 创建图形
    plt.figure(figsize=(12, 12))
    
    # 绘制样本
    for i, idx in enumerate(digit_indices[:num_samples]):
        # 获取图像和标签
        image, label = dataset[idx]
        
        # 将图像转换为numpy数组
        image = image.numpy()[0]  # 取第一个通道
        
        # 绘制子图
        plt.subplot(rows, cols, i + 1)
        plt.imshow(image, cmap='gray')
        plt.title(f"样本 #{i+1}")
        plt.axis('off')
    
    plt.suptitle(f"数字 {digit} 的样本", fontsize=16)
    plt.tight_layout()
    
    # 保存图像
    if save_path is None:
        save_path = f'digit_{digit}_samples.png'
    plt.savefig(save_path)
    print(f"数字 {digit} 的样本可视化已保存为 '{save_path}'")
    
    # 显示图像
    plt.show()

def main():
    # 加载数据集
    train_dataset, test_dataset = load_mnist_dataset()
    
    # 可视化随机样本
    visualize_mnist_samples(train_dataset)
    
    # 可视化每个数字的样本
    for digit in range(10):
        visualize_digit_samples(train_dataset, digit, num_samples=16, rows=4, cols=4)

if __name__ == "__main__":
    main() 